import org.ansj.domain.Result;
import org.ansj.domain.Term;
import org.ansj.splitWord.analysis.NlpAnalysis;
import org.ansj.splitWord.analysis.ToAnalysis;
import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDF;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Chinese words segmentation with user-dict in com.kingcall.dic
 * use Ansj(a java open source analyzer)
 */

// 这个信息就是你每次使用desc 进行获取函数信息的时候返回的
@Description(name = "ansj_seg", value = "_FUNC_(str) - chinese words segment using ansj. Return list of words.",
        extended = "Example: select _FUNC_('我是测试字符串') from src limit 1;\n"
                + "[\"我\", \"是\", \"测试\", \"字符串\"]")

public class NlpSeg extends UDF {

    public String evaluate(String input) {
        if (input.isEmpty() || "".equals(input)) {
            return input;
        }

        ArrayList<String> outResult = new ArrayList<>();

        //只关注这些词性的词
        Set<String> expectedNature = new HashSet<String>() {{
            add("n");add("v");add("vd");add("vn");add("vf");
            add("vx");add("vi");add("vl");add("vg");
            add("nt");add("nz");add("nw");add("nl");
            add("ng");add("userDefine");add("wh");add("nr");add("nr1");add("nr2");
            add("nrj");add("nrf");add("ns");add("nsf");
        }};
//        String str = "欢迎使用ansj_seg,(ansj中文分词)在这里如果你遇到什么问题都可以联系我.我一定尽我所能.帮助大家.ansj_seg更快,更准,更自由!";
        Result result = NlpAnalysis.parse(input); //分词结果的一个封装，主要是一个List<Term>的terms
//        System.out.println(result.getTerms());

        List<Term> terms = result.getTerms(); //拿到terms
//        System.out.println(terms.size());

        for (int i = 0; i < terms.size(); i++) {
            String word = terms.get(i).getName(); //拿到词
            String natureStr = terms.get(i).getNatureStr(); //拿到词性
            if (expectedNature.contains(natureStr)) {
//                System.out.println(word + ":" + natureStr);
                outResult.add(word);
            }
        }
        return outResult.size() > 0 ? String.join("|", outResult) : "";

    }
}

